package org.example.ws;

import com.sage.toolkit.gson.GsonUtils;
import lombok.SneakyThrows;
import lombok.extern.slf4j.Slf4j;
import org.springframework.boot.context.event.ApplicationStartedEvent;
import org.springframework.context.event.EventListener;
import org.springframework.web.socket.*;
import org.springframework.web.socket.handler.AbstractWebSocketHandler;

import java.io.IOException;
import java.util.Collection;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;

import static org.example.ws.StringConstants.*;

public abstract class AbsChatWebSocketHandler extends AbstractWebSocketHandler {
    // 保存每个房间的ChatRoom，一个ChatRoom包含多个用户session
    private static final Map<Object, ChatRoom> CHAT_ROOMS = new ConcurrentHashMap<>();
    // 保存每个用户的session
    private static final Map<Object, WebSocketSession> USER_SESSIONS = new ConcurrentHashMap<>();

    /**
     * 获取一个房间可加入的最大用户数
     *
     * @return
     */
    protected abstract int getMaxUserInRoom();

    /**
     * 建立连接后将当前session加入到聊天室，需要配合ChatSessionInterceptor使用
     *
     * @param session
     * @throws Exception
     */
    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        Object room = this.getRoom(session);
        ChatRoom chatRoom = CHAT_ROOMS.compute(room, (k, v) -> v == null ? new ChatRoom(room) : v);
        if (chatRoom.sessionCount() == this.getMaxUserInRoom()) {
            session.close(CloseStatus.POLICY_VIOLATION.withReason("房间用户数已达上限"));
            return;
        }

        chatRoom.addSession(session);
        Object user = this.getUser(session, true);
        USER_SESSIONS.put(user, session);
        this.afterRoomJoined(session);
    }

    /**
     * 加入聊天室后的自定义逻辑
     *
     * @param session
     */
    protected abstract void afterRoomJoined(WebSocketSession session);

    /**
     * 断开连接后将session从聊天室移除
     *
     * @param session
     * @param status
     * @throws Exception
     */
    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
        Object room = this.getRoom(session);
        ChatRoom chatRoom = CHAT_ROOMS.get(room);
        if (chatRoom != null) {
            chatRoom.removeSession(session);
        }

        Object user = this.getUser(session, true);
        USER_SESSIONS.remove(user);
        this.afterRoomLeaved(session);
    }

    /**
     * 离开聊天室后的自定义逻辑
     *
     * @param session
     */
    protected abstract void afterRoomLeaved(WebSocketSession session);

    /**
     * 处理文本消息，默认将消息广播给聊天室除了当前session外的所有session
     *
     * @param session
     * @param message
     * @throws Exception
     */
    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
        this.to(session).emit(message);
    }

    /**
     * 处理二进制消息，默认将消息广播给聊天室除了当前session外的所有session
     *
     * @param session
     * @param message
     * @throws Exception
     */
    @Override
    protected void handleBinaryMessage(WebSocketSession session, BinaryMessage message) throws Exception {
        this.to(session).emit(message);
    }

    /**
     * 获取当前session所在的房间
     *
     * @param session
     * @return
     */
    protected Object getRoom(WebSocketSession session) {
        return session.getAttributes().get(ROOM_KEY);
    }

    /**
     * 获取当前用户
     *
     * @param session
     * @param serialization
     * @return
     */
    protected Object getUser(WebSocketSession session, boolean serialization) {
        Object user = session.getAttributes().get(USER_KEY);
        return serialization ? GsonUtils.toJsonTree(user) : user;
    }

    protected Object getUser(WebSocketSession session) {
        return this.getUser(session, false);
    }

    /**
     * 获取指定房间的聊天室
     *
     * @param room
     * @return
     */
    public ChatRoom to(Object room) {
        return CHAT_ROOMS.getOrDefault(room, ChatRoom.EMPTY);
    }

    /**
     * 获取当前session所在的聊天室，后续所有操作都会忽略当前session
     *
     * @param session
     * @return
     */
    protected ChatRoom to(WebSocketSession session) {
        Object room = this.getRoom(session);
        return this.to(room).ignore(session);
    }

    /**
     * 向指定用户发送消息
     *
     * @param message
     * @param user
     * @param <T>
     */
    @SneakyThrows
    protected <T> void sendToUser(WebSocketMessage<T> message, Object user) {
        WebSocketSession session = USER_SESSIONS.get(user);
        if (session != null) {
            synchronized (session) {
                session.sendMessage(message);
            }
        }
    }

    @EventListener(ApplicationStartedEvent.class)
    public void onApplicationStartedEvent() {
        PingTask.runOnce(1000);
    }

    @Slf4j
    static class PingTask implements Runnable {
        private static final AtomicInteger counter = new AtomicInteger();
        private static final ScheduledExecutorService scheduler = Executors.newSingleThreadScheduledExecutor();

        static void runOnce(long pingFixedRate) {
            if (counter.incrementAndGet() == 1) {
                scheduler.scheduleAtFixedRate(new PingTask(), 0, pingFixedRate, TimeUnit.MILLISECONDS);
            }
        }

        @Override
        public void run() {
            Collection<WebSocketSession> sessions = USER_SESSIONS.values();
            for (WebSocketSession session : sessions) {
                try {
                    session.sendMessage(new PingMessage());
                } catch (IOException e) {
                    log.error(String.format("发送心跳异常. session = %s", session), e);
                }
            }
        }
    }
}
